function [deformedSrc optimCP optimAlpha] = pairRegistration(source, target, controlPointPrior)
% pairRegistration this function registers the source image to the target
% image...

%% to display debug info and plots
verbose = 0;

%% registration an gradient descent parameters
para.sigmaK = 3;
para.gamma = 0.1; % was 0.1
para.stepfactorCP = 0.1;
para.stepfactorAlpha = 0.1;
% sparsity parameters
para.gammasp = 0.000;
para.Csp = 10^(-1);

para.moveCtrlPts = 0;
para.maxIterations = 50;
para.maxTauIter = 10;
para.blurGamma = 0.1;

% termination criterion breakratio
para.breakRatio = 10^(-4);

% initialise control points on a regular grid and the momenta vectors to
% zero.
[nrow ncol] = size(source);
% sample on a regular 5x5 grid
%noGridPts = 15;
%[x y] = meshgrid(ncol/(noGridPts+1):ncol/(noGridPts+1):ncol-1,nrow/(noGridPts+1):nrow/(noGridPts+1):nrow-1);
%[x y] = meshgrid(para.sigmaK:para.sigmaK:ncol-1,para.sigmaK:para.sigmaK:nrow-1);

%[x y] = meshgrid(1:16);
% to be changed - need to wait for the results of the new data.
%x = [2 2 4 4 4 6 6 6 6 7 8 9 9 9 10 11 12 12 12 14 15 15]';
%y = [7 12 8 10 13 3 6 12 14 9 3 7 11 14 5 10 3 8 13 9 3 12]';
x = controlPointPrior(1,:)'; y = controlPointPrior(2,:)';

cp = [x(:), y(:)]';

alpha = zeros(size(cp));


%% preprocessing images
% if(strcmp(class(source),'uint8'))
%     source = double(source./255);
%     target = double(target./255);
% end
% 
source = gaussianBlur(source,para.blurGamma);
target = gaussianBlur(target,para.blurGamma);
source = gaussianBlur(source,para.blurGamma);
target = gaussianBlur(target,para.blurGamma);
%%
CP = zeros([size(cp),para.maxIterations]);
ALPHA = zeros([size(alpha),para.maxIterations]);
CP(:,:,1) = cp;
ALPHA(:,:,1) = alpha;
taucp = para.stepfactorCP;
taualpha = para.stepfactorAlpha;
[Enow iDiff regularizer sparsity] = Energy(CP(:,:,1),ALPHA(:,:,1),source,target,para);
Eprev = Enow;
E0 = Enow;
% perform the gradient descent
% [Eprev iDiff regularizer sparsity] = Energy(CP(:,:,1),ALPHA(:,:,1),source,target,para);
for iter = 1:para.maxIterations
    tauIter = 0;
    while((Eprev<=Enow)&&(tauIter<para.maxTauIter))        
        tauIter = tauIter + 1;
        GRADALPHA = gradAlpha(CP(:,:,iter),ALPHA(:,:,iter),source,target,para);
        ALPHA(:,:,iter+1) = ALPHA(:,:,iter) - taualpha.*GRADALPHA;
        if(para.moveCtrlPts==1)
            GRADCP = gradCP(CP(:,:,iter),ALPHA(:,:,iter),source,target,para);
            CP(:,:,iter+1) = CP(:,:,iter) - taucp.*GRADCP;
        else
            CP(:,:,iter+1) = CP(:,:,iter);
        end
        [Enow iDiff regularizer sparsity] = Energy(CP(:,:,iter+1),ALPHA(:,:,iter+1),source,target,para);
        taucp = taucp/2; taualpha = taualpha/2;
        if((tauIter >= para.maxTauIter)&&(Eprev<=Enow))
            ALPHA(:,:,iter+1) = ALPHA(:,:,iter);
            CP(:,:,iter+1) = CP(:,:,iter);
        end
    end
    % if we are confident about the stepsize then increase it by a factor
    % of 2 - 4 since we have already decreased it above
    taucp = 4*taucp; taualpha = 4*taualpha;
    if(verbose==1)
        disp(['Energy = ',num2str(Enow),' ImageDiff = ',num2str(iDiff),' Regularity = ',num2str(regularizer), ' Sparsity = ',num2str(sparsity)]);
    end
    % disp(['BreakRatio = ', num2str(((Eprev-Enow)/(E0-Enow)))]);
    % the following 4 variables are only for display of the energy graph
    Eprevtest(iter) = Enow;
    Difftest(iter) = iDiff;
    Regularizertest(iter) = regularizer;
    Sparsitytest(iter) = sparsity;
    
    % updating the values of the returned parameters
    deformedSrc = displayDeformedImage(CP(:,:,iter+1),ALPHA(:,:,iter+1), source, target, para);
    
    if(verbose == 1)
        % for displaying the deformed image at the current point    
        imagesc(deformedSrc); colormap gray;
        hold on;
        quiver(CP(1,:,iter+1), CP(2,:,iter+1), ALPHA(1,:,iter+1), ALPHA(2,:,iter+1));
        hold off;
        pause(0.01);
    end
    if(((Eprev-Enow)/(E0-Enow))<=para.breakRatio)
        break;
    end
    Eprev = Enow;    
end
% updating values of returned parameters
optimCP = CP(:,:,iter+1);
optimAlpha = ALPHA(:,:,iter+1);

Regularizertest = Regularizertest*para.gamma;
Sparsitytest = Sparsitytest*para.gammasp;

if(verbose == 1)
    figure; plot(Eprevtest,'Color','red'); hold on; 
    plot(Difftest,'Color','green'); 
    plot(Regularizertest,'Color','blue'); 
    %plot(Sparsitytest,'Color','cyan');
    %legend('Energy','imageDifference',['regularization with Gamma = ',num2str(para.gamma)], ...
    %['sparsity with GammaSP = ',num2str(para.gammasp),' C = ',num2str(para.Csp)],'Location','NorthEastOutside');
    legend('Energy','imageDifference',['regularization with Gamma = ',num2str(para.gamma)],'Location','NorthEastOutside');
end

deformedSrc = displayDeformedImage(optimCP,optimAlpha, source, target, para);
figure; subplot(1,3,1); imagesc(deformedSrc); title('deformed image'); colormap gray;
subplot(1,3,2); imagesc(source); title('source image'); colormap gray;
subplot(1,3,3); imagesc(target); title('target image'); colormap gray;

end

function I1 = displayDeformedImage(cp, alpha, Isrc, Itar, para)
    [X Y] = meshgrid(1:size(Isrc,2), 1:size(Isrc,1));
    yk = [X(:), Y(:)]';
    vyk = velocityField(yk,alpha,cp,para);
    Xvk = reshape(vyk(1,:),size(Isrc,1),size(Isrc,2));
    Yvk = reshape(vyk(2,:),size(Isrc,1),size(Isrc,2));
    % now is the fun part in which you get the new deformed source I1 which
    % is Isrc(yk-vyk)
    I1 = interp2(X,Y,Isrc,X-Xvk,Y-Yvk);
    I1(isnan(I1)) = 0;
end

function val = gaussianKer(c1, c2, para)
% this function is a simple gaussian kernel which takes the co-ordinates of
% 2 points as an input and returns the scalar which is the measure of the
% similarity between 2 points computed from the L2 norm of their difference    
    val = exp(-sum((c1-c2).^2)./(para.sigmaK^2));
end

function val = velocityField(x,alpha,cp,para)
val = zeros(2,size(x,2));
for i = 1:size(cp,2)
    val = val + repmat(gaussianKer(x,repmat(cp(:,i),1,size(x,2)),para),2,1).*repmat(alpha(:,i),1,size(x,2));
end
end

function val = gradAlpha(cp, alpha, Isrc, Itar, para)
    % first compute the matrix containing the position of the deformed grid
    % points from the target to the source. We are computing the term
    % yk - v(yk)
    % first define the regular grid
    [X Y] = meshgrid(1:size(Isrc,2), 1:size(Isrc,1));
    yk = [X(:), Y(:)]';
    vyk = velocityField(yk,alpha,cp,para);
    Xvk = reshape(vyk(1,:),size(Isrc,1),size(Isrc,2));
    Yvk = reshape(vyk(2,:),size(Isrc,1),size(Isrc,2));
    % now is the fun part in which you get the new deformed source I1 which
    % is Isrc(yk-vyk)
    I1 = interp2(X,Y,Isrc,X-Xvk,Y-Yvk);
    I1(isnan(I1)) = 0;
    [I1gradX I1gradY] = gradient(I1);
    Idiff = I1 - Itar;    
    Idiff = Idiff(:)'; I1grad = [I1gradX(:), I1gradY(:)]';
    val = zeros(2,size(alpha,2));
    % the following 2 loops are designed to run for all the control points
    for j = 1:size(alpha,2)
        % term2 = gaussianKer(cp, repmat(cp(:,j),1,size(cp,2)),para);
        term2 = alpha(:,j)*gaussianKer(cp, repmat(cp(:,j),1,size(cp,2)),para);
        % added the new sparsity penalty term
        if(sqrt(sum(alpha(:,j).^2))>=para.Csp)
            term3 = repmat(alpha(:,j)./sqrt(sum(alpha(:,j).^2)),1,size(alpha,2));
        else
            term3 = zeros(size(alpha));
        end
        val = val + para.gamma.*term2 + para.gammasp.*term3;
    end
    for k = 1:size(yk,2)
        %term1 = gaussianKer(cp, repmat(yk(:,k),1,size(cp,2)),para);
        term1 = I1grad(:,k)*(gaussianKer(cp, repmat(yk(:,k),1,size(cp,2)),para)*Idiff(:,k));
        val = val - term1;
    end
end

function val = gradCP(cp, alpha, Isrc, Itar, para)
    % first compute the matrix containing the position of the deformed grid
    % points from the target to the source. We are computing the term
    % yk - v(yk)
    % first define the regular grid
    [X Y] = meshgrid(1:size(Isrc,2), 1:size(Isrc,1));
    yk = [X(:), Y(:)]';
    vyk = velocityField(yk,alpha,cp,para);
    Xvk = reshape(vyk(1,:),size(Isrc,1),size(Isrc,2));
    Yvk = reshape(vyk(2,:),size(Isrc,1),size(Isrc,2));
    % now is the fun part in which you get the new deformed source I1 which
    % is Isrc(yk-vyk)
    I1 = interp2(X,Y,Isrc,X-Xvk,Y-Yvk);
    I1(isnan(I1)) = 0;
    [I1gradX I1gradY] = gradient(I1);
    Idiff = I1 - Itar;    
    Idiff = Idiff(:)'; I1grad = [I1gradX(:), I1gradY(:)]';
    val = zeros(2,size(cp,2));
    % the following 2 loops are designed to run for all the control points
    for k = 1:size(yk,2)
        term1 = gaussianKer(cp, repmat(yk(:,k),1,size(cp,2)),para);
        term1 = (I1grad(:,k)'*alpha).*(term1.*Idiff(:,k));
        term1 = repmat(term1,2,1).*(cp - repmat(yk(:,k),1,size(cp,2)));
        val = val + term1;
    end
    for j = 1:size(alpha,2)
        term2 = gaussianKer(cp, repmat(cp(:,j),1,size(cp,2)),para);
        term2 = ((alpha'*alpha(:,j))'.*term2);
        term2 = repmat(term2,2,1).*(cp - repmat(cp(:,j),1,size(cp,2)));
        val = val - para.gamma.*term2;
    end
    val = val.*(2/para.sigmaK^2);
end

function [val ImageDiff Regularizer Sparsity] = Energy(cp, alpha, Isrc, Itar, para)
    % first compute the matrix containing the position of the deformed grid
    % points from the target to the source. We are computing the term
    % yk - v(yk)
    % first define the regular grid
    [X Y] = meshgrid(1:size(Isrc,2), 1:size(Isrc,1));
    yk = [X(:), Y(:)]';
    vyk = velocityField(yk,alpha,cp,para);
    Xvk = reshape(vyk(1,:),size(Isrc,1),size(Isrc,2));
    Yvk = reshape(vyk(2,:),size(Isrc,1),size(Isrc,2));
    % now is the fun part in which you get the new deformed source I1 which
    % is Isrc(yk-vyk)
    I1 = interp2(X,Y,Isrc,X-Xvk,Y-Yvk);
    
    % the pixels outside the image all have a value of 0
    I1(isnan(I1)) = 0;
    
    Idiff = I1 - Itar;
    term1 = sum(Idiff(:).^2);
    val = zeros(1,size(alpha,2));
    term3 = 0;
    for j = 1:size(alpha,2)
        term2 = gaussianKer(cp, repmat(cp(:,j),1,size(cp,2)),para);
        term2 = ((alpha'*alpha(:,j))'.*term2);
        val = val + term2;
        term3 = term3 + LogC(sqrt(sum(alpha(:,j).^2)), para.Csp);
    end
    term2 = sum(val);    
    
    ImageDiff = term1;
    Regularizer = term2;
    Sparsity = term3;
    val = term1 + para.gamma*term2 + para.gammasp*term3;
end

function [val] = LogC(x,C)
    val = max(log(x), log(C)) - log(C);
end

function [G] = gaussianBlur(I,sigma)
H = zeros(size(I));
for i = 1:size(I,1)
    for j = 1:size(I,2)
        H(i,j) = exp(-(i^2+j^2)/sigma^2);
    end
end
H = fft2(H);
F = fft2(I,size(I,1),size(I,2));
G = real(ifft2(H.*F));
G = (G - min(G(:)))/(max(G(:)) - min(G(:)));
end